from chromadb import Embeddings
import requests
import json

# Ollama API 端点和API密钥

class QwenEmbedding(Embeddings):
    def __init__(self):
        self.ollamaApiUrl = "http://10.221.114.53:11434/v1/embeddings"  # 替换为实际的API端点
        self.apiKey = "your_api_key_here"  # 替换为你的API密钥

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return list(map(self.embed_query, texts))

    def embed_query(self, text: str) -> list[float]:
        # 设置请求头
        headers = {
            "Authorization": f"Bearer {self.apiKey}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": "qwen2:1.5b",  # 指定使用的模型版本
            "input": text
        }
        
        response = requests.post(self.ollamaApiUrl, headers=headers, data=json.dumps(payload))
        
        if response.status_code == 200:
            return response.json().get('data', [{}])[0].get('embedding')
        else:
            print(f"Error: Received status code {response.status_code}")
            print(f"Response body: {response.text}")
            return None
